import os
import pandas as pd
from flask import Flask, jsonify, send_file
import xml.etree.ElementTree as ET
import xmind
from xmind.core.markerref import MarkerId
from docx import Document

app = Flask(__name__)

def generate_directory_tree(root_folder):
    tree = {}
    for root, dirs, files in os.walk(root_folder):
        for file in files:
            if file.endswith('.xlsx'):
                # 分割文件名获取节点路径
                parts = file.replace('.xlsx', '').split('_')
                current_level = tree
                for part in parts:
                    if part not in current_level:
                        current_level[part] = {}
                    current_level = current_level[part]
                
                # 读取 Excel 文件获取指标名称，设置第二行为列名
                file_path = os.path.join(root, file)
                df = pd.read_excel(file_path, header=1)
                if '指标名称' in df.columns:
                    indicators = df['指标名称'].dropna().tolist()
                    current_level['指标'] = indicators
    
    return tree

def print_directory_tree(tree, level=0):
    for key, value in tree.items():
        print('  ' * level + key)
        if isinstance(value, dict):
            print_directory_tree(value, level + 1)
        elif isinstance(value, list):
            for item in value:
                print('  ' * (level + 1) + item)

def tree_to_excel(tree, file_path):
    data = []
    def traverse(node, path=''):
        for key, value in node.items():
            new_path = path + '/' + key if path else key
            if isinstance(value, dict):
                traverse(value, new_path)
            elif isinstance(value, list):
                for item in value:
                    data.append([new_path, item])
    traverse(tree)
    df = pd.DataFrame(data, columns=['路径', '指标'])
    df.to_excel(file_path, index=False)

def tree_to_xml(tree, file_path):
    root = ET.Element('root')
    
    def traverse(node, parent):
        for key, value in node.items():
            element = ET.SubElement(parent, 'node')
            element.set('name', key)
            if isinstance(value, dict):
                traverse(value, element)
            elif isinstance(value, list):
                for item in value:
                    sub_element = ET.SubElement(element, 'indicator')
                    sub_element.text = item
    
    traverse(tree, root)
    tree_xml = ET.ElementTree(root)
    tree_xml.write(file_path, encoding='utf-8', xml_declaration=True)

def tree_to_xmind(tree, file_path):
    try:
        workbook = xmind.load(file_path)
    except FileNotFoundError:
        workbook = xmind.create_workbook()
    
    sheet = workbook.getPrimarySheet()
    root_topic = sheet.getRootTopic()

    def traverse(node, parent_topic):
        for key, value in node.items():
            child_topic = parent_topic.addSubTopic()
            child_topic.setTitle(key)
            if isinstance(value, dict):
                traverse(value, child_topic)
            elif isinstance(value, list):
                note_content = '\n'.join(value)
                notes = child_topic.getNotes()
                if notes is None:
                    from xmind.core import notes as xmind_notes
                    from xmind.core import const
                    notes_element = xmind_notes.NotesElement()
                    # 修改设置笔记内容的方式
                    plain_notes = xmind_notes.NotesElement()
                    plain_notes_content = xmind_notes.NotesContent()
                    plain_notes_content.setTextContent(note_content)
                    plain_notes.setContent(plain_notes_content)
                    plain_notes.setType(const.PLAIN_NOTES)
                    notes_element.setElement(plain_notes)
                    child_topic.setNotes(notes_element)
                else:
                    # 若 notes 存在，暂时保持原逻辑，若报错需进一步调整
                    # 这里假设可以通过某种方式更新内容
                    # 实际可能需要重新设置整个 notes
                    pass

    traverse(tree, root_topic)
    xmind.save(workbook)

def tree_to_mm(tree, file_path):
    root = ET.Element('map')
    root.set('version', '1.0.1')

    def traverse(node, parent):
        for key, value in node.items():
            element = ET.SubElement(parent, 'node')
            element.set('TEXT', key)
            if isinstance(value, dict):
                traverse(value, element)
            elif isinstance(value, list):
                for item in value:
                    note = ET.SubElement(element, 'node')
                    note.set('TEXT', item)

    traverse(tree, root)
    tree_xml = ET.ElementTree(root)
    tree_xml.write(file_path, encoding='utf-8', xml_declaration=True)

def tree_to_docx(tree, file_path):
    doc = Document()
    def traverse(node, level=0):
        for key, value in node.items():
            # 根据层级添加标题
            if level <= 9:
                doc.add_heading(key, level)
            else:
                doc.add_paragraph('  ' * level + key)
            if isinstance(value, dict):
                traverse(value, level + 1)
            elif isinstance(value, list):
                for item in value:
                    doc.add_paragraph('  ' * (level + 1) + item)
    
    traverse(tree)
    doc.save(file_path)

def tree_to_markdown(tree, file_path):
    markdown_content = []
    def traverse(node, level=0):
        for key, value in node.items():
            # 根据层级添加标题
            heading = '#' * (level + 1)
            markdown_content.append(f"{heading} {key}")
            if isinstance(value, dict):
                traverse(value, level + 1)
            elif isinstance(value, list):
                for item in value:
                    markdown_content.append(f"{'  ' * (level + 1)}- {item}")
    
    traverse(tree)
    with open(file_path, 'w', encoding='utf-8') as f:
        f.write('\n'.join(markdown_content))

@app.route('/tree')
def get_tree():
    root_folder = 'root'  # 请将此处替换为实际的 root 文件夹路径
    tree = generate_directory_tree(root_folder)
    # 保存为 Excel 文件
    tree_to_excel(tree, 'tree_data.xlsx')
    # 保存为 XML 文件
    tree_to_xml(tree, 'tree_data.xml')
    # 保存为 XMind 文件
    # tree_to_xmind(tree, 'tree_data.xmind')
    # 保存为 Freemind 格式文件
    tree_to_mm(tree, 'tree_data.mm')
    # 保存为 DOCX 文件
    tree_to_docx(tree, 'tree_data.docx')
    # 保存为 Markdown 文件
    tree_to_markdown(tree, 'tree_data.md')
    return jsonify(tree)

@app.route('/')
def index():
    return send_file('index.html')

if __name__ == "__main__":
    # app.run(debug=True)
    # 这部分代码在启动 Flask 应用后不会执行，可移除
    root_folder = 'root'  
    tree = generate_directory_tree(root_folder)
    tree_to_markdown(tree, 'tree_data.md')
    print_directory_tree(tree)